{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "1f1c566e-412c-471c-8ee4-b1ab749cfbc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: [(1, 6), (1, 8), (9, 12), (9, 16)], 2: [(9, 14)]}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import collections\n",
    "import itertools\n",
    "\n",
    "\n",
    "def apply_pat_dont_know(current_pairs_set):\n",
    "    \"\"\"\n",
    "    Filters the current set of possible pairs based on Pat's statement\n",
    "    \"I don't know the numbers.\"\n",
    "    Pat knows the product P. If P can be formed by multiple pairs in\n",
    "    current_pairs_set, he doesn't know.\n",
    "    Returns a new set of pairs that are still possible.\n",
    "    \"\"\"\n",
    "    if not current_pairs_set:\n",
    "        return set()\n",
    "    \n",
    "    product_to_candidate_pairs = collections.defaultdict(list)\n",
    "    for x, y in current_pairs_set:\n",
    "        product_to_candidate_pairs[x * y].append((x, y))\n",
    "    \n",
    "    next_possible_pairs = set()\n",
    "    for product, candidates in product_to_candidate_pairs.items():\n",
    "        # If there's more than one pair for this product, Pat doesn't know.\n",
    "        # All such candidate pairs remain possible.\n",
    "        if len(candidates) > 1:\n",
    "            for pair in candidates:\n",
    "                next_possible_pairs.add(pair)\n",
    "    return next_possible_pairs\n",
    "\n",
    "def apply_sam_dont_know(current_pairs_set):\n",
    "    \"\"\"\n",
    "    Filters the current set of possible pairs based on Sam's statement\n",
    "    \"I don't know the numbers.\"\n",
    "    Sam knows the sum S. If S can be formed by multiple pairs in\n",
    "    current_pairs_set, she doesn't know.\n",
    "    Returns a new set of pairs that are still possible.\n",
    "    \"\"\"\n",
    "    if not current_pairs_set:\n",
    "        return set()\n",
    "        \n",
    "    sum_to_candidate_pairs = collections.defaultdict(list)\n",
    "    for x, y in current_pairs_set:\n",
    "        sum_to_candidate_pairs[x + y].append((x, y))\n",
    "        \n",
    "    next_possible_pairs = set()\n",
    "    for s, candidates in sum_to_candidate_pairs.items():\n",
    "        # If there's more than one pair for this sum, Sam doesn't know.\n",
    "        # All such candidate pairs remain possible.\n",
    "        if len(candidates) > 1:\n",
    "            for pair in candidates:\n",
    "                next_possible_pairs.add(pair)\n",
    "    return next_possible_pairs\n",
    "\n",
    "def find_solutions_pat_knows(current_pairs_set):\n",
    "    \"\"\"\n",
    "    Finds solutions after Pat's final statement \"I do know the numbers.\"\n",
    "    Pat knows the product P. If P uniquely identifies a pair in\n",
    "    current_pairs_set, that pair is a solution.\n",
    "    Returns a list of solution pairs.\n",
    "    \"\"\"\n",
    "    if not current_pairs_set:\n",
    "        return []\n",
    "        \n",
    "    product_to_candidate_pairs = collections.defaultdict(list)\n",
    "    for x, y in current_pairs_set:\n",
    "        product_to_candidate_pairs[x * y].append((x, y))\n",
    "        \n",
    "    solutions = []\n",
    "    for product, candidates in product_to_candidate_pairs.items():\n",
    "        # If there's exactly one pair for this product, Pat now knows.\n",
    "        if len(candidates) == 1:\n",
    "            solutions.append(candidates[0])\n",
    "    \n",
    "    solutions.sort() # For consistent output ordering\n",
    "    return solutions\n",
    "\n",
    "def solve_puzzle(numbers=range(1, 100), iterations=range(1, 10)):\n",
    "    \"\"\"\n",
    "    Solves the Sum and Product puzzle for various numbers of repetitions (n).\n",
    "    The dialogue is:\n",
    "    Pat: I don't know.\n",
    "    Sam: I don't know.\n",
    "    ... (repeated n times) ...\n",
    "    Pat: I do know the numbers.\n",
    "    Returns a dictionary where keys are n (number of repetitions) and\n",
    "    values are lists of solution pairs (x,y).\n",
    "    \"\"\"\n",
    "    all_solutions_by_n_repetitions = {}\n",
    "    \n",
    "    initial_candidate_pairs = list(itertools.combinations(numbers, 2))\n",
    "    \n",
    "    # Set a maximum for n repetitions to check.\n",
    "    # This can be adjusted; higher values take longer.\n",
    "    max_n_to_check = 12 # Adjusted for potentially faster feedback\n",
    "\n",
    "\n",
    "    for n_rep in range(1, max_n_to_check + 1):\n",
    "        current_possible_pairs = initial_candidate_pairs.copy()\n",
    "        \n",
    "        possible_to_reach_final_stage = True\n",
    "        for i_round in range(n_rep):\n",
    "            # Pat's turn: \"I don't know\"\n",
    "            current_possible_pairs = apply_pat_dont_know(current_possible_pairs)\n",
    "            # print(f\"  Round {i_round+1} (P): Pairs remaining = {len(current_possible_pairs)}\")\n",
    "            if not current_possible_pairs:\n",
    "                possible_to_reach_final_stage = False\n",
    "                break\n",
    "            \n",
    "            # Sam's turn: \"I don't know\"\n",
    "            current_possible_pairs = apply_sam_dont_know(current_possible_pairs)\n",
    "            # print(f\"  Round {i_round+1} (S): Pairs remaining = {len(current_possible_pairs)}\")\n",
    "            if not current_possible_pairs:\n",
    "                possible_to_reach_final_stage = False\n",
    "                break\n",
    "        \n",
    "        if not possible_to_reach_final_stage:\n",
    "            #print(f\"  Set of possible pairs became empty before Pat's final statement for n = {n_rep}.\")\n",
    "            continue \n",
    "\n",
    "        if not current_possible_pairs:\n",
    "            #print(f\"  Set of possible pairs is empty before Pat's final statement for n = {n_rep}.\")\n",
    "            continue\n",
    "\n",
    "        #print(f\"  Pairs remaining before Pat's final statement for n = {n_rep}: {len(current_possible_pairs)}\")\n",
    "        \n",
    "        solutions_for_this_n = find_solutions_pat_knows(current_possible_pairs)\n",
    "        \n",
    "        if solutions_for_this_n:\n",
    "            all_solutions_by_n_repetitions[n_rep] = solutions_for_this_n\n",
    "            \n",
    "    return all_solutions_by_n_repetitions\n",
    "\n",
    "solve_puzzle(numbers=range(1, 20))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "d7e40017-85d8-49b2-a36d-c557d0b3d10c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{2: {(1, 6), (1, 8), (72, 92), (75, 96)},\n",
       " 3: {(81, 88)},\n",
       " 4: {(77, 90)},\n",
       " 5: {(76, 90)},\n",
       " 6: {(80, 84)},\n",
       " 7: {(77, 84)},\n",
       " 8: set(),\n",
       " 9: set(),\n",
       " 10: set(),\n",
       " 11: set()}"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "63a43e89-550e-48f4-8582-e7bcc5d062e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: set(),\n",
       " 2: set(),\n",
       " 3: set(),\n",
       " 4: set(),\n",
       " 5: set(),\n",
       " 6: set(),\n",
       " 7: set(),\n",
       " 8: set(),\n",
       " 9: set()}"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve({(x, y) for x in range(2, 100) for y in range(x + 1, 100) if x + y <= 100})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "66e2c711-9a46-41b1-abcc-b1d53b50657b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(1, 6),\n",
       " (1, 8),\n",
       " (2, 3),\n",
       " (2, 4),\n",
       " (2, 6),\n",
       " (2, 9),\n",
       " (3, 4),\n",
       " (3, 6),\n",
       " (3, 8),\n",
       " (4, 6)}"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pairs = set(combinations(range(1, 10), 2))\n",
    "apply_pat_dont_know(pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "31725942-75a5-4613-9d8d-f40127903b87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(1, 6),\n",
       " (1, 8),\n",
       " (2, 3),\n",
       " (2, 4),\n",
       " (2, 6),\n",
       " (2, 9),\n",
       " (3, 4),\n",
       " (3, 6),\n",
       " (3, 8),\n",
       " (4, 6)}"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unknown_pairs(prod, pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "79cdfef6-c87f-4456-814b-1cd8e74e18df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: [(1, 6), (1, 8), (72, 92), (75, 96)],\n",
       " 2: [(81, 88)],\n",
       " 3: [(77, 90)],\n",
       " 4: [(76, 90)],\n",
       " 5: [(80, 84)],\n",
       " 6: [(77, 84)]}"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve_puzzle(numbers=range(1, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "5092bc62-7c82-4831-9038-83be1c25344c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: {(1, 6), (1, 8), (72, 92), (75, 96)},\n",
       " 2: {(81, 88)},\n",
       " 3: {(77, 90)},\n",
       " 4: {(76, 90)},\n",
       " 5: {(80, 84)},\n",
       " 6: {(77, 84)},\n",
       " 7: set(),\n",
       " 8: set(),\n",
       " 9: set()}"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve(numbers=range(1, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51bc1771-ed13-4c97-8dee-5de1d1e51bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#solve_puzzle(numbers=range(1, 10))\n",
    "numbers=range(1, 10)\n",
    "pairs = sum_pairs = prod_pairs = set(combinations(numbers, 2))\n",
    "sum_pairs = unknown_pairs(prod, sum_pairs)\n",
    "prod_pairs = unknown_pairs(sum, prod_pairs)\n",
    "sum_pairs, prod_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "eb6b40ec-9d62-4d66-8eb9-5d6c95f4822d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(1, 2), (5, 2)}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "known_pairs(prod, {(1, 2), (2, 3), (6, 1), (3, 4), (6, 2), (5, 2)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "78df4550-e213-4b36-aa57-ee993502cd9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(1, 2), (5, 2)}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "\n",
    "known_pairs(prod, {(1, 2), (2, 3), (6, 1), (3, 4), (6, 2), (5, 2)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b61e2589-bad6-48be-9b83-f0f1b28c8c60",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
